#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <std_msgs/Int32.h>
#include <mavros_msgs/AttitudeTarget.h>
#include <tf/tf.h>

// 全局变量
mavros_msgs::State current_state;
geometry_msgs::PoseStamped local_pos;
int target_number = -1;
int target_index = -1;
double pose_x = 0;
double pose_y = 0;
double pose_z = 0;
bool rotation_done = false;
ros::Publisher attitude_pub;
double initial_yaw = 0;

// 回调函数
void stateCallback(const mavros_msgs::State::ConstPtr& msg) {
    current_state = *msg;
}

void localPosCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
    local_pos = *msg;
}

void numberCallback(const std_msgs::Int32::ConstPtr& msg) {
    target_number = msg->data;
}

// 根据 target_index 更新目标位置
void updatePoseBasedOnIndex() {
    if (target_number == -1) return;
	
     target_index = target_number;
    if (target_index == 3 || target_index == 4 || target_index == 9 || target_index == 10 ||
        target_index == 13 || target_index == 18 || target_index == 19 || target_index == 24) {
        pose_x = 1.75;
    } else if (target_index == 2 || target_index == 5 || target_index == 8 || target_index == 11 ||
               target_index == 14 || target_index == 17 || target_index == 20 || target_index == 23) {
        pose_x = 1.25;
    } else {
        pose_x = 0.75;
    }

    int mode_y = (target_index-1) / 6;
    switch (mode_y) {
        case 0: pose_y = 0; break;
        case 1: pose_y = -1.25; break;
        case 2: pose_y = -2.25; break;
        case 3: pose_y = -3.5; break;
        default: pose_y = 0; break;
    }

    if ((target_index >= 4 && target_index <= 9) || (target_index >= 16 && target_index <= 21)) {
        pose_z = 1;
    } else {
        pose_z = 1.4;
    }
}

// 从四元数获取偏航角
double getYawFromQuaternion(const geometry_msgs::Quaternion& quat) {
    tf::Quaternion tf_quat;
    tf::quaternionMsgToTF(quat, tf_quat);
    double roll, pitch, yaw;
    tf::Matrix3x3(tf_quat).getRPY(roll, pitch, yaw);
    return yaw;
}

// 执行步骤
void executeStep(int& step, int& sametimes, geometry_msgs::PoseStamped& pose, ros::Publisher& local_pos_pub, ros::ServiceClient& set_mode_client, ros::ServiceClient& arming_client, ros::Time& last_request) {
    mavros_msgs::SetMode offb_set_mode; // 移动到函数开头
    switch (step) {
        case 0: // 初始上升
            pose.pose.position.x = 0;
            pose.pose.position.y = 0;
            pose.pose.position.z = pose_z;
            if (fabs(local_pos.pose.position.z - pose_z) < 0.1) {
                if (sametimes > 20) {
                    if (pose_y == -1.25 || pose_y == -3.5) {
                        
                        initial_yaw = getYawFromQuaternion(local_pos.pose.orientation); // 获取初始偏航角
                        double target_yaw = initial_yaw + M_PI;

                        tf::Quaternion q;
                        q.setRPY(0, 0, target_yaw);
                        geometry_msgs::Quaternion q_msg;
                        tf::quaternionTFToMsg(q, q_msg);

                        pose.pose.orientation = q_msg;

                        // 发布旋转姿态
                        local_pos_pub.publish(pose);
                            step = 1; // 完成自转后继续执行后续步骤
                            pose.pose.position.x = -0.4;
                            pose.pose.position.y = 0; // 设置目标y位置
                            pose.pose.position.z = pose_z;
                        
                    } else {
                        step = 1; // 不需要自转，直接进入下一步
                        pose.pose.position.x = -0.4;           
                        pose.pose.position.y = 0;
                        pose.pose.position.z = pose_z;
                    }
                } else {
                    sametimes++;
                }
            } else {
                sametimes = 0;
            }
            break;
	case 1: // 前往目标y位置
            if (fabs(local_pos.pose.position.x+0.4)<0.1){
                if (sametimes > 20) {
                    step = 2;
                    pose.pose.position.x = -0.4;
                    pose.pose.position.y = pose_y;
                    pose.pose.position.z = pose_z;
                } else {
                    sametimes++;
                }
            } else {
                sametimes = 0;
            }
            break;

        case 2: // 前往目标y位置
            if (fabs(local_pos.pose.position.y - pose_y) < 0.1) {
                if (sametimes > 20) {
                    step = 3;
                    pose.pose.position.x = pose_x;
                    pose.pose.position.y = pose_y;
                    pose.pose.position.z = pose_z;
                } else {
                    sametimes++;
                }
            } else {
                sametimes = 0;
            }
            break;

        case 3: // 前往目标x位置
            if (fabs(local_pos.pose.position.x - pose_x) < 0.1) {
                if (sametimes > 40) {
                    step = 4;
                    pose.pose.position.x = 2.6;
                    pose.pose.position.y = pose_y;
                    pose.pose.position.z = pose_z;
                } else {
                    sametimes++;
                }
            } else {
                sametimes = 0;
            }
            break;

        case 4: // 前往固定x位置
            if (local_pos.pose.position.x > 2.5 && local_pos.pose.position.x < 2.7){
                if (sametimes > 20) {
                    step = 5;
                    pose.pose.position.x = 2.6;
                    pose.pose.position.y = -3.5;
                    pose.pose.position.z = pose_z;
                } else {
                    sametimes++;
                }
            } else {
                sametimes = 0;
            }
            break;
	case 5: // 前往固定x位置
            if (local_pos.pose.position.y > -3.6&& local_pos.pose.position.y < -3.4){
                if (sametimes > 20) {
                    step = 6;
                    pose.pose.position.x = 2.4;
                    pose.pose.position.y = -3.5;
                    pose.pose.position.z = pose_y;
                } else {
                    sametimes++;
                }
            } else {
                sametimes = 0;
            }
            break;

        case 6: // 前往固定y位置
            if (local_pos.pose.position.x> 2.3 && local_pos.pose.position.x < 2.5) {
                if (sametimes > 20) {
                    step = 6;
                } else {
                    sametimes++;
                }
            } else {
                sametimes = 0;
            }
            break;

        case 7: // 准备降落
            sametimes = 0;
            offb_set_mode.request.custom_mode = "AUTO.LAND";
            if (current_state.mode != "AUTO.LAND" && (ros::Time::now() - last_request > ros::Duration(5.0))) {
                if (set_mode_client.call(offb_set_mode) && offb_set_mode.response.mode_sent) {
                    ROS_INFO("AUTO.LAND enabled");
                }
                last_request = ros::Time::now();
        }

            break;
    }
}

int main(int argc, char **argv) {
    ros::init(argc, argv, "offb_cfx");
    ros::NodeHandle nh;

    // 订阅者和发布者
    ros::Subscriber state_sub = nh.subscribe("mavros/state", 10, stateCallback);
    ros::Subscriber local_pos_sub = nh.subscribe("mavros/local_position/pose", 10, localPosCallback);
    ros::Subscriber number_sub = nh.subscribe("target_number", 10, numberCallback);
    ros::Publisher local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>("mavros/setpoint_position/local", 10);
    attitude_pub = nh.advertise<mavros_msgs::AttitudeTarget>("mavros/setpoint_attitude/attitude", 10);

    // 服务客户端
    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>("mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>("mavros/set_mode");

    // 发布频率
    ros::Rate rate(20.0);

    // 等待与FCU连接
    while (ros::ok() && !current_state.connected) {
        ros::spinOnce();
        rate.sleep();
    }

    geometry_msgs::PoseStamped pose;
    pose.pose.position.x = 0;
    pose.pose.position.y = 0;
    pose.pose.position.z = 1.5;

    // 发布一些设定点
    for (int i = 0; ros::ok() && i < 100; ++i) {
        local_pos_pub.publish(pose);
        ros::spinOnce();
        rate.sleep();
    }

    // 设置无人机模式和解锁
    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";
    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;

    ros::Time last_request = ros::Time::now();
    int step = 0;
    int sametimes = 0;

    while (ros::ok()) {
        // 无人机状态设定与判断
        if (current_state.mode != "OFFBOARD" && (ros::Time::now() - last_request > ros::Duration(5.0))) {
            if (set_mode_client.call(offb_set_mode) && offb_set_mode.response.mode_sent) {
                ROS_INFO("Offboard enabled");
            }
            last_request = ros::Time::now();
        } else if (!current_state.armed && (ros::Time::now() - last_request > ros::Duration(5.0))) {
            if (arming_client.call(arm_cmd) && arm_cmd.response.success) {
                ROS_INFO("Vehicle armed");
            }
            last_request = ros::Time::now();
        } else {
            updatePoseBasedOnIndex();  // 更新pose_x, pose_y, pose_z
            if (target_number != -1) {
                ROS_INFO("Position: x %f, y %f, z %f", pose_x, pose_y, pose_z);
                executeStep(step, sametimes, pose, local_pos_pub, set_mode_client, arming_client, last_request);
                // 发布位置控制信息
                local_pos_pub.publish(pose);
            }
        }

        ros::spinOnce();
        rate.sleep(); // 控制循环频率
    }

    return 0;
}


